import os
import sys
from Bio import SeqIO
import pyhmmer

"""
使用方法：python collect_by_hmmer.py hmm_file.fasta fasta_folder n

结果筛选：在 run_hmmer 函数中，根据 n 是否为 None 来决定是否对结果进行筛选。如果 n 不为 None，则按 E-value 排序并筛选前 n 个结果；否则使用所有结果。

参数:
- hmm_file: HMM模型文件路径。
- fasta_folder: 包含FASTA格式序列文件的文件夹路径。
- n (可选): 筛选的最佳结果数量，默认为不筛选。
- output_folder: 存储搜索结果和匹配序列的输出文件夹路径。

fasta文件存放在fasta_proteomes/example/，hmm文件存放在hmm_profiles/all_lectin/，使用hmmer,将hmm文件与fasta_proteomes/example/中的所有fasta文件比对，将比对结果存放在比对结果文件夹mining_result中，比对结果文件名为hmm文件名+fasta文件名，并且取出所有匹配的蛋白全序列，保存为fasta文件。
"""

def run_hmmer(hmm_file, fasta_folder, output_folder, n=None):
    # 确保输出文件夹存在，如果不存在则创建
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    # 打开HMM文件并读取模型
    with pyhmmer.plan7.HMMFile(hmm_file) as hmm_file:
        hmm = hmm_file.read()

    # 遍历FASTA文件夹，处理每个FASTA文件
    for fasta_file in os.listdir(fasta_folder):
        if fasta_file.endswith('.fasta') or fasta_file.endswith('.fa'):
            fasta_path = os.path.join(fasta_folder, fasta_file)
            result_file = os.path.join(output_folder, f"{os.path.basename(hmm_file)}_{fasta_file}.out")
            matched_sequences_file = os.path.join(output_folder, f"{os.path.basename(hmm_file)}_{fasta_file}_matched.fasta")

            # 读取FASTA文件中的序列
            with pyhmmer.easel.SequenceFile(fasta_path) as seq_file:
                sequences = list(seq_file)

            # 使用HMM模型搜索序列
            with pyhmmer.hmmsearch.Pipeline(alphabet=hmm.alphabet) as pipeline:
                results = pipeline.search_hmm(hmm, sequences)

            # 获取所有结果
            all_hits = results[0]

            # 如果指定了n，则筛选前n个结果
            if n is not None:
                top_results = sorted(all_hits, key=lambda x: x.evalue)[:n]
            else:
                top_results = all_hits

            # 写入搜索结果
            with open(result_file, 'w') as out_handle:
                for hit in top_results:
                    out_handle.write(f"Hit: {hit.name}\n")
                    out_handle.write(f"E-value: {hit.evalue}\n")
                    out_handle.write(f"Score: {hit.score}\n")
                    out_handle.write(f"Bias: {hit.bias}\n")
                    out_handle.write("\n")

            # 收集匹配的序列
            matched_sequences = []
            for hit in top_results:
                for domain in hit.domains:
                    sequence = sequences[hit.idx]
                    matched_sequences.append(sequence)

            # 如果有匹配的序列，则写入匹配序列文件
            if matched_sequences:
                with open(matched_sequences_file, 'w') as matched_handle:
                    SeqIO.write(matched_sequences, matched_handle, 'fasta')

if __name__ == "__main__":
    if len(sys.argv) < 3:
        print("Usage: python collect_by_hmmer.py hmm_file fasta_folder [n]")
        sys.exit(1)

    hmm_file = sys.argv[1]
    fasta_folder = sys.argv[2]
    output_folder = "mining_result"
    n = int(sys.argv[3]) if len(sys.argv) > 3 else None

    run_hmmer(hmm_file, fasta_folder, output_folder, n)